#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <std_msgs/Int32MultiArray.h>
#include <mavros_msgs/CommandBool.h>
#include <mavros_msgs/SetMode.h>
#include <mavros_msgs/State.h>
#include <std_msgs/Bool.h>
#include <vector>
#include <cmath>
#include <iomanip>

mavros_msgs::State current_state;
geometry_msgs::PoseStamped current_pose;
std::vector<geometry_msgs::PoseStamped> waypoints;
bool path_received = false;
bool hovering = false;
bool descending = false;
ros::Time hover_start_time;
int waypoint_idx = 0;
std_msgs::Bool hover_msg;
geometry_msgs::PoseStamped target_pose;
geometry_msgs::PoseStamped descent_start_pose;
geometry_msgs::PoseStamped descent_goal_pose;

geometry_msgs::PoseStamped mapped_current_pose; // 用于保存并复用“上一个航点”
double dx, dy, dz;

void state_cb(const mavros_msgs::State::ConstPtr& msg) {
    current_state = *msg;
}

void pose_cb(const geometry_msgs::PoseStamped::ConstPtr& msg) {
    current_pose = *msg;
}

void path_cb(const std_msgs::Int32MultiArray::ConstPtr& msg) {
    if (path_received) return;

    const auto& data = msg->data;
    int n = data.size();
    if (n < 6) return;

    int i = 0;
    while (i < n - 1) {
        if (data[i] == -2 && data[i+1] == -2) {
            i += 2;
            continue;
        }
        if (data[i] == -9 && data[i+1] == -9) break;

        geometry_msgs::PoseStamped pose;
        pose.header.frame_id = "map";
        pose.pose.position.x = data[i] * 0.5;
        pose.pose.position.y = data[i+1] * 0.5;
        pose.pose.position.z = 1.2;
        waypoints.push_back(pose);
        i += 2;
    }
    path_received = true;
    ROS_INFO("Received path with %ld waypoints.", waypoints.size());
}

bool reached_waypoint(const geometry_msgs::PoseStamped& wp) {
    double dx = wp.pose.position.x - current_pose.pose.position.x;
    double dy = wp.pose.position.y - current_pose.pose.position.y;
    double dz = wp.pose.position.z - current_pose.pose.position.z;
    return std::sqrt(dx*dx + dy*dy + dz*dz) < 0.3;
}

int main(int argc, char **argv) {
    ros::init(argc, argv, "offb_2025final");
    ros::NodeHandle nh;

    ros::Subscriber state_sub = nh.subscribe("iris_0/mavros/state", 10, state_cb);
    ros::Subscriber pose_sub = nh.subscribe("iris_0/mavros/local_position/pose", 10, pose_cb);
    ros::Subscriber path_sub = nh.subscribe("path_topic", 10, path_cb);
    ros::Publisher local_pos_pub = nh.advertise<geometry_msgs::PoseStamped>("iris_0/mavros/setpoint_position/local", 10);
    // 发布映射后“上一个航点”的话题
    ros::Publisher mapped_target_pub = nh.advertise<geometry_msgs::PoseStamped>("mapped_coordinate_point", 10);
    ros::Publisher hover_pub = nh.advertise<std_msgs::Bool>("is_hovering", 10);

    ros::ServiceClient arming_client = nh.serviceClient<mavros_msgs::CommandBool>("iris_0/mavros/cmd/arming");
    ros::ServiceClient set_mode_client = nh.serviceClient<mavros_msgs::SetMode>("iris_0/mavros/set_mode");
    
    ros::Rate rate(20.0);

    while (ros::ok() && !current_state.connected) {
        ros::spinOnce();
        rate.sleep();
    }

    target_pose.pose.position.x = 0;
    target_pose.pose.position.y = 0;
    target_pose.pose.position.z = 1.2;
    target_pose.pose.orientation.w = 1.0;

    for (int i = 0; i < 100; ++i) {
        local_pos_pub.publish(target_pose);
        ros::spinOnce();
        rate.sleep();
    }

    mavros_msgs::SetMode offb_set_mode;
    offb_set_mode.request.custom_mode = "OFFBOARD";

    mavros_msgs::CommandBool arm_cmd;
    arm_cmd.request.value = true;

    ros::Time last_request = ros::Time::now();

    while (ros::ok()) {
        if (current_state.mode != "OFFBOARD" &&
            (ros::Time::now() - last_request > ros::Duration(5.0))) {
            if (set_mode_client.call(offb_set_mode) && offb_set_mode.response.mode_sent) {
                ROS_INFO("Offboard enabled");
            }
            last_request = ros::Time::now();
        }

        if (!current_state.armed &&
            (ros::Time::now() - last_request > ros::Duration(5.0))) {
            if (arming_client.call(arm_cmd) && arm_cmd.response.success) {
                ROS_INFO("Vehicle armed");
            }
            last_request = ros::Time::now();
        }

        if (path_received && waypoint_idx < waypoints.size()) {
            // 发布“上一个”航点（如果存在）
            if (waypoint_idx > 0) {
                mapped_current_pose = waypoints[waypoint_idx - 1];
                mapped_current_pose.header.stamp = ros::Time::now();
                mapped_current_pose.header.frame_id = "map";
                mapped_target_pub.publish(mapped_current_pose);
            }

            target_pose = waypoints[waypoint_idx];
            local_pos_pub.publish(target_pose);
            
            ros::Duration(0.5).sleep();
            ROS_INFO_STREAM("Step to WP " << waypoint_idx << ": ("
                << std::fixed << std::setprecision(2)
                << target_pose.pose.position.x << ", "
                << target_pose.pose.position.y << ", "
                << target_pose.pose.position.z << ")");

            if (reached_waypoint(target_pose)) {
                waypoint_idx++;

                if (waypoint_idx == waypoints.size() - 3) {
                    hover_start_time = ros::Time::now();
                    hovering = true;
                    ROS_INFO("Reached third-to-last waypoint, start hovering...");
                }
            }
        }
        else if (hovering) {
            if ((ros::Time::now() - hover_start_time).toSec() < 4.0) {
                target_pose = waypoints[waypoints.size() - 3];
                local_pos_pub.publish(target_pose);
                hover_msg.data = true;
                hover_pub.publish(hover_msg);
                
                ros::Duration(0.5).sleep();
            } else {
                hovering = false;
                descending = true;
                descent_start_pose = waypoints[waypoints.size() - 3];
                target_pose = descent_start_pose;

                double delta_z = descent_start_pose.pose.position.z - 0.2;
                descent_goal_pose.pose.position.z = 0.2;

                if (std::abs(descent_start_pose.pose.position.x) >= std::abs(descent_start_pose.pose.position.y)) {
                    descent_goal_pose.pose.position.x = descent_start_pose.pose.position.x > 0 ?
                        descent_start_pose.pose.position.x - delta_z :
                        descent_start_pose.pose.position.x + delta_z;
                    descent_goal_pose.pose.position.y = descent_start_pose.pose.position.y;
                } else {
                    descent_goal_pose.pose.position.y = descent_start_pose.pose.position.y > 0 ?
                        descent_start_pose.pose.position.y - delta_z :
                        descent_start_pose.pose.position.y + delta_z;
                    descent_goal_pose.pose.position.x = descent_start_pose.pose.position.x;
                }

                dx = descent_goal_pose.pose.position.x - descent_start_pose.pose.position.x;
                dy = descent_goal_pose.pose.position.y - descent_start_pose.pose.position.y;
                dz = descent_goal_pose.pose.position.z - descent_start_pose.pose.position.z;
                double norm = std::sqrt(dx*dx + dy*dy + dz*dz);
                if (norm < 1e-6) norm = 1.0;
                dx /= norm; dy /= norm; dz /= norm;
                hover_msg.data = true;
                hover_pub.publish(hover_msg);
                ROS_INFO("Start Afghan-style angled descent...");
            }
        }
        else if (descending) {
            double step = 0.05;
            target_pose.pose.position.x += dx * step;
            target_pose.pose.position.y += dy * step;
            target_pose.pose.position.z += dz * step;
            ROS_INFO("dx=%.2f dy=%.2f dz=%.2f", dx, dy, dz);
            local_pos_pub.publish(target_pose);
            hover_msg.data = true;
            hover_pub.publish(hover_msg);
            ros::Duration(0.5).sleep();
            double dist_to_goal = std::sqrt(
                std::pow(target_pose.pose.position.x - descent_goal_pose.pose.position.x, 2) +
                std::pow(target_pose.pose.position.y - descent_goal_pose.pose.position.y, 2) +
                std::pow(target_pose.pose.position.z - descent_goal_pose.pose.position.z, 2)
            );

            if (dist_to_goal <= 0.05) {
                mavros_msgs::SetMode land_mode;
                land_mode.request.custom_mode = "AUTO.LAND";
                set_mode_client.call(land_mode);
                ROS_INFO("Reached goal, switching to AUTO.LAND");
                break;
            }
        }

        ros::spinOnce();
        rate.sleep();
    }

    return 0;
}

